from typing import List


class Solution:

    def findDiagonalOrder(self, matrix: List[List[int]]) -> object:
        n = len(matrix)
        m = len(matrix[0])
        p = [[[-1, 1], [1, -1]], [[0, 1], [1, 0]]]
        ans = []
        x, y = 0, 0
        total = 0
        d1 = 0
        d2 = 0
        while total < n * m:
            ans.append(matrix[x][y])
            total += 1
            if total >= n * m:
                break
            xx = x + p[0][d1][0]
            yy = y + p[0][d1][1]
            if xx < 0 or xx >= n or yy < 0 or yy >= m:
                d2 = d1
                d1 = (d1 + 1) % 2
                xx = x + p[1][d2][0]
                yy = y + p[1][d2][1]
                while xx < 0 or xx >= n or yy < 0 or yy >= m:
                    d2 = (d2 + 1) % 2
                    xx = x + p[1][d2][0]
                    yy = y + p[1][d2][1]
            x, y = xx, yy
        return ans


mm = [
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9]
]

solve = Solution()
print(mm)
print(solve.findDiagonalOrder(mm))
